import requests
import json

from web.scraper.line2_upload_image import Line2UploadImage


class AllAiLine2:
    def __init__(self):
        self.domain = "https://yuanbao.tencent.com"  # 腾讯元宝
        self.my_key = "naQivTmsDa" # 改
        self.headers = {
            'authority': 'yuanbao.tencent.com',
            'x-agentid': f'{self.my_key}',
            'chat_version': 'v1',
            'x-requested-with': 'XMLHttpRequest',
            'x-source': 'web',
            'user-agent': 'Mozilla/5.0 (iPhone; CPU iPhone OS 13_2_3 like Mac OS X) AppleWebKit/605.1.15 (KHTML, '
                          'like Gecko)'
                          'Version/13.0.3 Mobile/15E148 Safari/604.1 Edg/92.0.4515.131',
            'content-type': 'text/plain;charset=UTF-8',
            'accept': '*/*',
            'origin': 'https://yuanbao.tencent.com',
            'sec-fetch-site': 'same-origin',
            'sec-fetch-mode': 'cors',
            'sec-fetch-dest': 'empty',
            'referer': f'https://yuanbao.tencent.com/chat/{self.my_key}',
            'accept-language': 'zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6',
        }
        self.cookies = {
            'hy_source': 'web',
            'hy_user': 'c903525ab38242efa670fc443640557b', # 改
            'hy_token': 'amlqbRV7Feb+gXpvj/ixQsAX1rY/82hWXh1L0v4fBL68JhuorWmZxyjAJaiDUQ9iobrSkCmdjt5Sga7dzDZGJA==', # 改
        }
        self.upload_image = Line2UploadImage(self.cookies, self.headers) # 上传图片类

    # 创建对话
    def create_dialogue(self):
        json_data = {
            'agentId': self.my_key,
        }

        response = requests.post(f'{self.domain}/api/user/agent/conversation/create', cookies=self.cookies,
                                 headers=self.headers, json=json_data).json()
        return response

    # 开始会话
    def start_dialogue(self, dialogue_id, message='你叫什么名字', stream=False):
        data = json.dumps({"model": "gpt_175B_0404", "prompt": message, "plugin": "Adaptive", "displayPrompt": message,
                          "displayPromptType": 1, "options": {
                "imageIntention": {"needIntentionModel": True, "backendUpdateFlag": 2, "intentionStatus": False}},
                          "multimedia": [], "agentId": self.my_key, "supportHint": 1, "version": "v2"})
        response = requests.post(f'https://yuanbao.tencent.com/api/chat/{dialogue_id}', cookies=self.cookies,
                                 headers=self.headers, data=data, stream=True)
        if stream:
            return response
        else:
            content = ""

            for line in response.iter_lines():
                if line:
                    try:
                        text = line.decode('utf-8')
                        if text.startswith('data: '):
                            text = text.replace('data: ', '')
                            if text == "[DONE]":
                                break
                            text = json.loads(text)
                            content += text['msg']

                    except json.JSONDecodeError as e:
                        pass
                    except KeyError as e:
                        pass
            return content

    # 获取对话内容
    def get_dialogue_content(self, dialogue_id):
        json_data = {
            'conversationId': dialogue_id,
            'offset': 0,
            'limit': 10,
        }

        response = requests.post('https://yuanbao.tencent.com/api/user/agent/conversation/v1/detail',
                                 cookies=self.cookies, headers=self.headers, json=json_data)
        return response.json()

    # 获取对话列表
    def get_dialogue_list(self, limit=1):
        json_data = {
            'agentId': self.my_key,
            'offset': 0,
            'limit': limit,
            'filterGoodQuestion': True,
        }

        response = requests.post('https://yuanbao.tencent.com/api/user/agent/conversation/list',
                                 cookies=self.cookies, headers=self.headers, json=json_data)

        return response.json()

    # 翻译流(内部函数)
    def translate_stream(self, response):
        content = ""
        for line in response.iter_lines():
            if line:
                try:
                    text = line.decode('utf-8')
                    if text.startswith('data: '):
                        text = text.replace('data: ', '')
                        if text == "[DONE]":
                            break
                        text = json.loads(text)
                        print(text['msg'], end="", flush=True)
                        content += text['msg']


                except json.JSONDecodeError as e:
                    pass
                except KeyError as e:
                    pass
        return content


    # 识别图片内容
    def identify_image_content(self, image_url, user_message=None, stream=False, dialogue_id=None):
        if dialogue_id is None:
            return {
                "code": 404,
                "message": "对话ID不存在"
            }
        if user_message is None:
            user_message = "请精确提取图片中的文字内容，以纯文本形式呈现，不要使用 Markdown 语法、代码或任何其他干扰性内容。严格按照以下格式回复：#IMAGETEXTSTART 你所提取到的图片全部文字内容 #IMAGETEXTEND 。确保提取的文字内容直接紧跟在 #IMAGETEXTSTART 之后，且在 #IMAGETEXTEND 之前，中间不能有 “图片中的文字内容” 等任何其他无关文字或表述。只输出符合此格式要求的内容，不要有任何额外说明。 "
        data = {
        "model":"gpt_175B_0404",
        "prompt":user_message,
        "plugin":"Adaptive",
        "displayPrompt":user_message,
        "displayPromptType":1,
        "options":{
            "imageIntention":{
                "needIntentionModel":True,
                "backendUpdateFlag":2,
                "intentionStatus":True
                }
            },
            "multimedia":[
                {
                    "type":"image",
                    "docType":"image",
                    "url":f"{image_url}",
                    "fileName":"",
                    "size":0,
                    "width":0,
                    "height":0
                }
            ],
            "agentId":self.my_key,
            "supportHint":1,
            "version":"v2",
            "chatModelId":"hunyuan_gpt_175B_0404"
        }
        data = json.dumps(data)
        response = requests.post(f'https://yuanbao.tencent.com/api/chat/{dialogue_id}', stream=True,cookies=self.cookies, headers=self.headers, data=data)
        return self.translate_stream_v1(response, stream)
    

    # 翻译流v1(内部函数)
    def translate_stream_v1(self, response, stream ):
        if stream:
            return response
        else:
            content = ""
            for line in response.iter_lines():
                if line:
                    try:
                        text = line.decode('utf-8')
                        if text.startswith('data: '):
                            text = text.replace('data: ', '')
                            if text == "[DONE]":
                                break
                            text = json.loads(text)
                            # print(text['msg'], end="", flush=True)
                            content += text['msg']


                    except json.JSONDecodeError as e:
                        pass
                    except KeyError as e:
                        pass
            return content
                    
    
if __name__ == "__main__":
    line1 = AllAiLine2()
    print(line1.identify_image_content('https://hunyuan.tencent.com/api/resource/download?resourceId=952b645dd081ac16f760e66bdaa1dae9', dialogue_id='a36134d3-aa74-4e66-a579-80726412ac7f'))
